from io import BytesIO
import pandas as pd
from fastapi import APIRouter, WebSocket, Depends, Body
from starlette.responses import StreamingResponse

from app.schemas.simulation import SimulationCreate, SimulationInDB, SimulationWithResults, SimulationQuery, PageResult, \
    SimulationWithResultsDetail
from app.crud import simulation as crud_sim
from typing import List

from app.utils.Result import Result

router = APIRouter()

@router.post(
    "/",
    response_model=Result[SimulationInDB],
    summary="新增 Simulation",
    description="创建一个新的仿真记录（Simulation），指定环境参数 ID 和可选的名称。"
)
async def create_simulation(simulation: SimulationCreate):
    new_sim = await crud_sim.create_simulation(simulation)
    return Result.ok(data=new_sim)

@router.get(
    "/",
    response_model=Result[List[SimulationInDB]],
    summary="查询所有 Simulation",
    description="获取所有仿真记录，按创建时间倒序排序返回。"
)
async def list_simulations():
    simulations = await crud_sim.get_all_simulations()
    return Result.ok(data=simulations)

@router.post(
    "/list-page",
    response_model=Result[PageResult],
    summary="分页查询 Simulation",
    description="分页获取仿真记录，支持关键词模糊匹配，按创建时间倒序排序。"
)
async def list_simulations_page(query: SimulationQuery = Body(...)):
    skip = (query.page - 1) * query.page_size
    simulations, total = await crud_sim.get_simulations_by_page(query.keyword, skip, query.page_size)

    page_data = PageResult(
        items=simulations,
        total=total,
        page=query.page,
        page_size=query.page_size
    )
    return Result.ok(data=page_data)


@router.get(
    "/with-results",
    response_model=Result[List[SimulationWithResults]],
    summary="查询所有仿真及其结果",
    description="获取所有 Simulation，并附带每条记录对应的仿真结果总表（SimulationResultTotal）。"
)
async def list_simulations_with_results():
    simulations = await crud_sim.get_all_simulations_with_results()
    return Result.ok(data=simulations)

@router.get(
    "/with-results/{sim_id}",
    response_model=Result[SimulationWithResults],
    summary="根据ID查询仿真及其结果",
    description="根据仿真ID获取单个 Simulation，并附带每条记录对应的仿真结果总表（SimulationResultTotal）。"
)
async def get_simulations_with_results(sim_id: int):
    simulations = await crud_sim.get_simulations_with_results(sim_id)
    return Result.ok(data=simulations)

@router.get(
    "/with-results-detail/{sim_id}",
    response_model=Result[SimulationWithResultsDetail],
    summary="根据ID查询仿真及其结果",
    description="根据仿真ID获取单个 Simulation，并附带对应的仿真结果总表（SimulationResultTotal及对应的SimulationResult）。"
)
async def get_simulation_with_results_detail(sim_id: int):
    simulation = await crud_sim.get_simulation_with_results(sim_id)
    if not simulation:
        return Result.error(msg="未找到指定ID的仿真记录")
    return Result.ok(data=simulation)

@router.delete(
    "/{simulation_id}",
    response_model=Result[str],
    summary="根据 ID 删除 Simulation",
    description="根据 simulation_id 级联删除指定仿真记录，若不存在则返回失败提示。"
)
async def delete_simulation(simulation_id: int):
    deleted = await crud_sim.delete_simulation(simulation_id)
    if deleted == 0:
        return Result.error(code=500, msg="删除失败")
    return Result.ok(msg="删除成功")

@router.websocket(
    "/start_simulation",
)
async def start_simulation(ws:WebSocket,env_param_id: int, simulation_id: int):
    await ws.accept()
    await crud_sim.run_simulation_rl(ws,env_param_id, simulation_id)
    return {"message": "Simulation started"}


@router.get(
    "/simulation/export_simulation_by_id/{sim_id}",
    summary="根据ID列表导出仿真结果",
    description="传入simulationID，导出对应Excel文件"
)
async def export_simulation_by_id_api(sim_id: int):
    # 获取仿真对象及其关联结果
    simulation_data = await crud_sim.get_simulation_with_results(sim_id)

    # 通过对象属性访问
    simulation_name = simulation_data.name
    results = simulation_data.results  # 假设是一个列表，每个元素是一个策略结果对象

    # 创建内存中的 Excel writer（不写入磁盘）
    output = BytesIO()
    writer = pd.ExcelWriter(output, engine='openpyxl')
    rows_simulation_total = []

    for result in results:
        strategy_name = result.strategy
        rows_simulation_total.append({
            '电池电量': result.battery_energy,
            '策略名称': strategy_name,
            '总供电量': result.total_power,
            '光伏供电量': result.solar_power,
            '风力供电量': result.wind_power,
            '国网供电量': result.grid_power,
            '光伏花费': result.solar_cost,
            '风力花费': result.wind_cost,
            '国网花费': result.grid_cost,
            '光伏碳排放量': result.solar_emission,
            '风力碳排放量': result.wind_emission,
            '国网碳排放量': result.grid_emission
        })
        result_maps = result.result_maps  # 假设也是对象列表

        rows = []
        for map_item in result_maps:
            r = map_item.simulation_result  # 访问 simulation_result 对象
            time_point = str(r.time_point)

            rows.append({
                '时间点': time_point,
                '电池电量': r.battery_energy,
                '太阳能': r.solar_power,
                '风能': r.wind_power,
                '国网': r.grid_power,
                '总供电量': r.total_power,
                '太阳费用': r.solar_cost,
                '风费用': r.wind_cost,
                '国网费用': r.grid_cost,
                '太阳碳排': r.solar_emission,
                '风碳排': r.wind_emission,
                '国网碳排': r.grid_emission,
            })

        # 转换为 DataFrame 并写入 Excel
        df = pd.DataFrame(rows)
        sheet_name = strategy_name[:31]  # Excel 限制 sheet 名最多31字符
        df.to_excel(writer, sheet_name=sheet_name, index=False)

    df = pd.DataFrame(rows_simulation_total)
    sheet_name = simulation_name[:31]  # Excel 限制 sheet 名最多31字符
    df.to_excel(writer, sheet_name=sheet_name, index=False)

    writer.close()
    output.seek(0)

    # 构造 StreamingResponse 返回 Excel 文件
    filename = f"simulation.xlsx"
    headers = {
        "Content-Disposition": f"attachment; filename*=UTF-8''{filename}"
    }
    return StreamingResponse(output, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", headers=headers)
